package org.example.common.loadbalancer.cloud;

import lombok.extern.slf4j.Slf4j;
import org.example.common.core.constants.GrayConstant;
import org.example.common.core.utils.ribbon.GrayRequestContextHandler;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.reactive.DefaultResponse;
import org.springframework.cloud.client.loadbalancer.reactive.EmptyResponse;
import org.springframework.cloud.client.loadbalancer.reactive.Request;
import org.springframework.cloud.client.loadbalancer.reactive.Response;
import org.springframework.cloud.loadbalancer.core.NoopServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.SelectedInstanceCallback;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import reactor.core.publisher.Mono;

import java.util.ArrayList;
import java.util.List;
import java.util.Random;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 自定义负载均衡规则（灰度发布）
 */
@Slf4j
public class GrayCustomLoadBalancerRule implements ReactorServiceInstanceLoadBalancer {
    private final String serviceId;
    private final AtomicInteger position;
    // 服务列表
    private final ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider;

    public GrayCustomLoadBalancerRule(String serviceId,
                            ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider) {
        this(serviceId, new Random().nextInt(1000), serviceInstanceListSupplierProvider);
    }

    public GrayCustomLoadBalancerRule(String serviceId, int seed,
                            ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider) {
        this.serviceId = serviceId;
        this.position = new AtomicInteger(seed);
        this.serviceInstanceListSupplierProvider = serviceInstanceListSupplierProvider;
    }

    @SuppressWarnings("rawtypes")
    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        try {
            boolean grayTag = GrayRequestContextHandler.getGrayTag();
            ServiceInstanceListSupplier supplier = serviceInstanceListSupplierProvider
                    .getIfAvailable(NoopServiceInstanceListSupplier::new);
            return supplier.get().next()
                    .map(serviceInstances -> processInstanceResponse(request, supplier, serviceInstances,grayTag));
        } finally {
            //清除灰度标记
            GrayRequestContextHandler.remove();
      }
    }

    private Response<ServiceInstance> processInstanceResponse(Request<?> request, ServiceInstanceListSupplier supplier,
                                                              List<ServiceInstance> serviceInstances,boolean grayTag) {
        //灰度发布的服务
        List<ServiceInstance> grayServerList = new ArrayList<>();
        //正常的服务
        List<ServiceInstance> normalServerList = new ArrayList<>();
        for(ServiceInstance server : serviceInstances) {
            //从nacos中获取元素剧进行匹配
            if(server.getMetadata().containsKey(GrayConstant.GRAY_HEADER)
                    && server.getMetadata().get(GrayConstant.GRAY_HEADER).equals(GrayConstant.GRAY_VALUE)) {
                grayServerList.add(server);
            } else {
                normalServerList.add(server);
            }
        }
        Response<ServiceInstance> serviceInstanceResponse =null;
        //如果被标记为灰度发布，则调用灰度发布的服务
        if(grayTag&&grayServerList.size()>0) {
            serviceInstanceResponse= getInstanceResponse(grayServerList);
        } else {
            serviceInstanceResponse= getInstanceResponse(normalServerList);
        }
        if (supplier instanceof SelectedInstanceCallback && serviceInstanceResponse.hasServer()) {
            ((SelectedInstanceCallback) supplier).selectedServiceInstance(serviceInstanceResponse.getServer());
        }
        return serviceInstanceResponse;
    }

    private Response<ServiceInstance> getInstanceResponse(List<ServiceInstance> instances) {
        if (instances.isEmpty()) {
            if (log.isWarnEnabled()) {
                log.warn("No servers available for service: " + serviceId);
            }
            return new EmptyResponse();
        }
        // 默认RoundRobinLoadBalancer
//         return this.getRoundRobinInstance(instances);
        // 随机
        return this.getRandomInstance(instances);
    }

    /**
     * 使用随机算法 参考{link
     * {@link org.springframework.cloud.loadbalancer.core.RandomLoadBalancer}}
     *
     * @param instances 实例
     * @return {@link Response }<{@link ServiceInstance }>
     */
    private Response<ServiceInstance> getRandomInstance(List<ServiceInstance> instances) {
        int index = ThreadLocalRandom.current().nextInt(instances.size());
        ServiceInstance instance = instances.get(index);
        return new DefaultResponse(instance);
    }

    /**
     * 使用RoundRobin机制获取节点
     *
     * @param instances 实例
     * @return {@link Response }<{@link ServiceInstance }>
     */
    private Response<ServiceInstance> getRoundRobinInstance(List<ServiceInstance> instances) {
        // 每一次计数器都自动+1，实现轮询的效果
        int pos = this.position.incrementAndGet() & Integer.MAX_VALUE;
        ServiceInstance instance = instances.get(pos % instances.size());
        return new DefaultResponse(instance);
    }
}
